# Author: Mark Richardson
# Date started: 01/28/2022
# Date finished: 02/07/2022
# Purpose: Format performance ratings for export
# Revised 10/20/2022
# Revision purpose: Format hierarchical ratings and rater parameters
# Proofed 11/05/2023

# Load packages
library(dplyr)
library(ggplot2)
library(rstan)
library(stringr)

# Load data
load("data/ratings/performance_ratings/02_performance_fitted_models.RData")

# Function to scale ratings to be N(0, 1)

localID <- function(stan_object) {
  
  x <- as.matrix(stan_object, pars = "x")
  
  x <- (x - mean(x)) / sd(x)
  
  x_out <- tibble(mn = apply(x, MARGIN = 2, FUN = mean),
                  sd = apply(x, MARGIN = 2, FUN = sd),
                  lb = apply(x, MARGIN = 2, FUN = quantile, probs = 0.025),
                  ub = apply(x, MARGIN = 2, FUN = quantile, probs = 0.975),
                  agency_index = 1:ncol(x))
  
  return(x_out)
    
}

#### Hierarchical performance ratings with informed priors using std. dev. ####

x_sd_hier <- localID(perf_inf_sd_hier) %>%
  full_join(., agency_data, by = "agency_index") %>%
  left_join(., perf_stan %>% 
              select(agency_index, agency_rated, n_ratings_agency) %>%
              distinct(), by = c("agency_index", "agency_rated")) %>%
  select(agency_rated, mn:ub, n_inf_ratings, n_ratings_agency) %>%
  rename(agency = agency_rated,
         perf_rating_inf_hier = mn,
         perf_rating_inf_hier_sd = sd,
         perf_rating_inf_hier_lb = lb,
         perf_rating_inf_hier_ub = ub,
         n_inf_perf_ratings = n_inf_ratings,
         n_perf_ratings = n_ratings_agency)

#### Performance ratings with informed priors using std. dev. ####

x_sd <- localID(perf_inf_sd) %>%
  full_join(., agency_data, by = "agency_index") %>%
  left_join(., perf_stan %>% 
              select(agency_index, agency_rated, n_ratings_agency) %>%
              distinct(), by = c("agency_index", "agency_rated")) %>%
  select(agency_rated, mn:ub, n_inf_ratings, n_ratings_agency) %>%
  rename(agency = agency_rated,
         perf_rating_inf = mn,
         perf_rating_inf_sd = sd,
         perf_rating_inf_lb = lb,
         perf_rating_inf_ub = ub,
         n_inf_perf_ratings = n_inf_ratings,
         n_perf_ratings = n_ratings_agency)

#### Hierarchical performance ratings with naive priors ####

x_nv_hier <- localID(perf_naive_sd_hier) %>%
  full_join(., agency_data, by = "agency_index") %>%
  left_join(., perf_stan %>% 
              select(agency_index, agency_rated, n_ratings_agency) %>%
              distinct(), by = c("agency_index", "agency_rated")) %>%
  select(agency_rated, mn:ub, n_inf_ratings, n_ratings_agency) %>%
  rename(agency = agency_rated,
         perf_rating_nv_hier = mn,
         perf_rating_nv_hier_sd = sd,
         perf_rating_nv_hier_lb = lb,
         perf_rating_nv_hier_ub = ub,
         n_inf_perf_ratings = n_inf_ratings,
         n_perf_ratings = n_ratings_agency)

#### Performance ratings with naive priors ####

x_nv <- localID(perf_naive) %>%
  full_join(., agency_data, by = "agency_index") %>%
  left_join(., perf_stan %>% 
              select(agency_index, agency_rated, n_ratings_agency) %>%
              distinct(), by = c("agency_index", "agency_rated")) %>%
  select(agency_rated, mn:ub, n_inf_ratings, n_ratings_agency) %>%
  rename(agency = agency_rated,
         perf_rating_nv = mn,
         perf_rating_nv_sd = sd,
         perf_rating_nv_lb = lb,
         perf_rating_nv_ub = ub,
         n_inf_perf_ratings = n_inf_ratings,
         n_perf_ratings = n_ratings_agency)

x <- full_join(x_sd, x_nv)

x <- full_join(x, x_sd_hier)

x <- full_join(x, x_nv_hier)

#### Save formatted data ####

rm(list = ls()[ls() != "x"])

perf_ratings <- x %>% 
  select(agency,
         perf_rating_inf_hier:perf_rating_inf_hier_ub,
         perf_rating_inf:perf_rating_inf_ub,
         perf_rating_nv_hier:perf_rating_nv_hier_ub,
         perf_rating_nv:perf_rating_nv_ub,
         n_inf_perf_ratings, n_perf_ratings) %>%
  mutate(n_inf_perf_ratings = if_else(is.na(n_inf_perf_ratings), 0, as.double(n_inf_perf_ratings)))

# Require 5 ratings per agency

perf_ratings <- perf_ratings %>% filter(n_perf_ratings >= 5)

save(perf_ratings, file = "data/ratings/performance_ratings/04_performance_ratings_fmtd.RData")

haven::write_dta(perf_ratings, path =  "data/ratings/performance_ratings/04_performance_ratings_fmtd.dta")

#### Format rater parameters ####

rm(list = ls())

# Load data
load("data/ratings/performance_ratings/02_performance_fitted_models.RData")

# Format betas

beta_inf_hier <- summary(perf_inf_sd_hier, pars = "beta")$summary %>%
  as_tibble(., rownames = "par") %>%
  mutate(rater_index = str_remove(par, "beta\\["),
         rater_index = str_remove(rater_index, "]"),
         rater_index = as.integer(rater_index)) %>%
  full_join(., rater_data, by = "rater_index") %>%
  full_join(., perf_stan %>% select(rater_index, id, dept_index, dept) %>% distinct()) %>%
  select(mean, sd, `2.5%`, `97.5%`, rater_index:dept) %>%
  rename(beta_inf_hier = mean,
         beta_inf_hier_sd = sd,
         beta_inf_hier_lb = `2.5%`,
         beta_inf_hier_ub = `97.5%`)

beta_inf <- summary(perf_inf_sd, pars = "beta")$summary %>%
  as_tibble(., rownames = "par") %>%
  mutate(rater_index = str_remove(par, "beta\\["),
         rater_index = str_remove(rater_index, "]"),
         rater_index = as.integer(rater_index)) %>%
  full_join(., rater_data, by = "rater_index") %>%
  select(mean, sd, `2.5%`, `97.5%`, rater_index:id) %>%
  rename(beta_inf = mean,
         beta_inf_sd = sd,
         beta_inf_lb = `2.5%`,
         beta_inf_ub = `97.5%`)

beta_nv_hier <- summary(perf_naive_sd_hier, pars = "beta")$summary %>%
  as_tibble(., rownames = "par") %>%
  mutate(rater_index = str_remove(par, "beta\\["),
         rater_index = str_remove(rater_index, "]"),
         rater_index = as.integer(rater_index)) %>%
  full_join(., rater_data, by = "rater_index") %>%
  select(mean, sd, `2.5%`, `97.5%`, rater_index:id) %>%
  rename(beta_nv_hier = mean,
         beta_nv_hier_sd = sd,
         beta_nv_hier_lb = `2.5%`,
         beta_nv_hier_ub = `97.5%`)

beta_nv <- summary(perf_naive, pars = "beta")$summary %>%
  as_tibble(., rownames = "par") %>%
  mutate(rater_index = str_remove(par, "beta\\["),
         rater_index = str_remove(rater_index, "]"),
         rater_index = as.integer(rater_index)) %>%
  full_join(., rater_data, by = "rater_index") %>%
  select(mean, sd, `2.5%`, `97.5%`, rater_index:id) %>%
  rename(beta_nv = mean,
         beta_nv_sd = sd,
         beta_nv_lb = `2.5%`,
         beta_nv_ub = `97.5%`)

beta <- full_join(beta_inf_hier, beta_inf) %>%
  full_join(., beta_nv_hier) %>%
  full_join(., beta_nv) %>%
  select(id, rater_index, dept, dept_index, everything())

# Format alphas

alpha_inf_hier <- summary(perf_inf_sd_hier, pars = "alpha")$summary %>%
  as_tibble(., rownames = "par") %>%
  mutate(rater_index = str_remove(par, "alpha\\["),
         rater_index = str_remove(rater_index, "]"),
         rater_index = as.integer(rater_index)) %>%
  full_join(., rater_data, by = "rater_index") %>%
  full_join(., perf_stan %>% select(rater_index, id, dept_index, dept) %>% distinct()) %>%
  select(mean, sd, `2.5%`, `97.5%`, rater_index:dept) %>%
  rename(alpha_inf_hier = mean,
         alpha_inf_hier_sd = sd,
         alpha_inf_hier_lb = `2.5%`,
         alpha_inf_hier_ub = `97.5%`)

alpha_nv_hier <- summary(perf_naive_sd_hier, pars = "alpha")$summary %>%
  as_tibble(., rownames = "par") %>%
  mutate(rater_index = str_remove(par, "alpha\\["),
         rater_index = str_remove(rater_index, "]"),
         rater_index = as.integer(rater_index)) %>%
  full_join(., rater_data, by = "rater_index") %>%
  select(mean, sd, `2.5%`, `97.5%`, rater_index:id) %>%
  rename(alpha_nv_hier = mean,
         alpha_nv_hier_sd = sd,
         alpha_nv_hier_lb = `2.5%`,
         alpha_nv_hier_ub = `97.5%`)

# Format non-hierarchical models - need to add mu_alpha

# Informed model
mu_alpha_inf <- as.matrix(perf_inf_sd, pars = "mu_alpha")

alpha_inf <- as.matrix(perf_inf_sd, pars = "alpha")

alpha_inf_all <- matrix(nrow = nrow(alpha_inf),
                        ncol = ncol(alpha_inf),
                        dimnames = dimnames(alpha_inf))

# Add mu_alpha to every draw of alpha

for (i in 1:ncol(alpha_inf)) {
  
  alpha_inf_all[,i] <- alpha_inf[,i] + mu_alpha_inf
  
}

rm(mu_alpha_inf, alpha_inf)

# Naive model
mu_alpha_nv <- as.matrix(perf_naive, pars = "mu_alpha")

alpha_nv <- as.matrix(perf_naive, pars = "alpha")

alpha_nv_all <- matrix(nrow = nrow(alpha_nv),
                        ncol = ncol(alpha_nv),
                        dimnames = dimnames(alpha_nv))

# Add mu_alpha to every draw of alpha

for (i in 1:ncol(alpha_nv)) {
  
  alpha_nv_all[,i] <- alpha_nv[,i] + mu_alpha_nv
  
}

rm(mu_alpha_nv, alpha_nv)

# Aggregate the draws to the rater level

alpha_inf <- tibble(alpha_inf = apply(alpha_inf_all, MARGIN = 2, FUN = mean),
                    alpha_inf_sd = apply(alpha_inf_all, MARGIN = 2, FUN = sd),
                    alpha_inf_lb = apply(alpha_inf_all, MARGIN = 2, FUN = quantile, probs = 0.025),
                    alpha_inf_ub = apply(alpha_inf_all, MARGIN = 2, FUN = quantile, probs = 0.975),
                    rater_index = 1:ncol(alpha_inf_all)) %>%
  full_join(., rater_data)

alpha_nv <- tibble(alpha_nv = apply(alpha_nv_all, MARGIN = 2, FUN = mean),
                   alpha_nv_sd = apply(alpha_nv_all, MARGIN = 2, FUN = sd),
                   alpha_nv_lb = apply(alpha_nv_all, MARGIN = 2, FUN = quantile, probs = 0.025),
                   alpha_nv_ub = apply(alpha_nv_all, MARGIN = 2, FUN = quantile, probs = 0.975),
                   rater_index = 1:ncol(alpha_nv_all)) %>%
  full_join(., rater_data)

alpha <- full_join(alpha_inf_hier, alpha_inf) %>%
  full_join(., alpha_nv_hier) %>%
  full_join(., alpha_nv) %>%
  select(id, rater_index, dept, dept_index, everything())

# Get dept hyperparameters for the hierarchical models

# Beta
mu_beta_inf <- summary(perf_inf_sd_hier, pars = "mu_beta")$summary %>%
  as_tibble(., rownames = "par") %>%
  mutate(dept_index = str_remove(par, "mu_beta\\["),
         dept_index = str_remove(dept_index, "]"),
         dept_index = as.integer(dept_index)) %>%
  full_join(., dept_data, by = "dept_index") %>%
  select(mean, sd, `2.5%`, `97.5%`, dept_index:n_raters_dept) %>%
  rename(mu_beta_inf = mean,
         mu_beta_inf_sd = sd,
         mu_beta_inf_lb = `2.5%`,
         mu_beta_inf_ub = `97.5%`)

mu_beta_nv <- summary(perf_naive_sd_hier, pars = "mu_beta")$summary %>%
  as_tibble(., rownames = "par") %>%
  mutate(dept_index = str_remove(par, "mu_beta\\["),
         dept_index = str_remove(dept_index, "]"),
         dept_index = as.integer(dept_index)) %>%
  full_join(., dept_data, by = "dept_index") %>%
  select(mean, sd, `2.5%`, `97.5%`, dept_index:n_raters_dept) %>%
  rename(mu_beta_nv = mean,
         mu_beta_nv_sd = sd,
         mu_beta_nv_lb = `2.5%`,
         mu_beta_nv_ub = `97.5%`)

beta <- full_join(beta, mu_beta_inf) %>%
  full_join(., mu_beta_nv)

# Alpha
mu_alpha_inf <- summary(perf_inf_sd_hier, pars = "mu_alpha")$summary %>%
  as_tibble(., rownames = "par") %>%
  mutate(dept_index = str_remove(par, "mu_alpha\\["),
         dept_index = str_remove(dept_index, "]"),
         dept_index = as.integer(dept_index)) %>%
  full_join(., dept_data, by = "dept_index") %>%
  select(mean, sd, `2.5%`, `97.5%`, dept_index:n_raters_dept) %>%
  rename(mu_alpha_inf = mean,
         mu_alpha_inf_sd = sd,
         mu_alpha_inf_lb = `2.5%`,
         mu_alpha_inf_ub = `97.5%`)

mu_alpha_nv <- summary(perf_naive_sd_hier, pars = "mu_alpha")$summary %>%
  as_tibble(., rownames = "par") %>%
  mutate(dept_index = str_remove(par, "mu_alpha\\["),
         dept_index = str_remove(dept_index, "]"),
         dept_index = as.integer(dept_index)) %>%
  full_join(., dept_data, by = "dept_index") %>%
  select(mean, sd, `2.5%`, `97.5%`, dept_index:n_raters_dept) %>%
  rename(mu_alpha_nv = mean,
         mu_alpha_nv_sd = sd,
         mu_alpha_nv_lb = `2.5%`,
         mu_alpha_nv_ub = `97.5%`)

alpha <- full_join(alpha, mu_alpha_inf) %>%
  full_join(., mu_alpha_nv)

# Save the rater parameters

save(alpha, beta, file = "data/ratings/performance_ratings/06_performance_rater_parameters_fmtd.RData")
